package com.yyfax.pay.common.filter;

import com.google.common.collect.Maps;
import com.yyfax.framework.common.LogConst;
import com.yyfax.framework.filter.YYFilter;
import com.yyfax.pay.common.util.UUIDTool;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.*;

/**
 * 追踪traceid
 * @author wangjwy
 * @date 2018-05-23 10:56:44
 * @since v1.0.0
 */
public class TraceFilter extends YYFilter {

	public static final String TRACEID = "traceId";
	public static final String TRACEID2 = "Trace-Id";
	@Override
	protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain)
			throws ServletException, IOException {
		try {

			ModifyParametersWrapper mParametersWrapper = new ModifyParametersWrapper(request);

			String uid = request.getHeader(TRACEID2);
			if (StringUtils.isBlank(uid)) {
				 uid = request.getHeader(TRACEID);
			}
			if (StringUtils.isBlank(uid)) {
				uid = UUIDTool.getUUID();
				mParametersWrapper.putHeader(TRACEID2, uid);
			}
			MDC.put(LogConst.URL, request.getServletPath());
			MDC.put(TRACEID, uid);
			response.addHeader(TRACEID, uid);
			filterChain.doFilter(mParametersWrapper, response);
		} finally {
			MDC.clear();
		}
	}

	/**
	 * 继承HttpServletRequestWrapper，创建装饰类，以达到修改HttpServletRequest参数的目的
	 */
	private class ModifyParametersWrapper extends HttpServletRequestWrapper {
		private final Map<String, String> customHeaders;

		ModifyParametersWrapper(HttpServletRequest request) {
			super(request);
			this.customHeaders = Maps.newHashMap();
		}

		void putHeader(String name, String value) {
			this.customHeaders.put(name, value);
		}

		@Override
		public String getHeader(String name) {
			// check the custom headers first
			String headerValue = customHeaders.get(name);

			if (headerValue != null) {
				return headerValue;
			}
			// else return from into the original wrapped object
			return ((HttpServletRequest) getRequest()).getHeader(name);
		}

		@Override
		public Enumeration<String> getHeaderNames() {
			// create a set of the custom header names
			Set<String> set = new HashSet<>(customHeaders.keySet());

			// now add the headers from the wrapped request object
			Enumeration<String> e = ((HttpServletRequest) getRequest()).getHeaderNames();
			while (e.hasMoreElements()) {
				// add the names of the request headers into the list
				String n = e.nextElement();
				set.add(n);
			}

			// create an enumeration from the set and return
			return Collections.enumeration(set);
		}
	}
}
